package cn.fetosoft.woodpecker.core.jmeter.processor;

import org.apache.commons.codec.binary.Base64;
import org.apache.commons.lang3.StringUtils;
import org.apache.jmeter.processor.PostProcessor;
import org.apache.jmeter.samplers.SampleResult;
import org.apache.jmeter.testelement.AbstractTestElement;
import org.apache.jmeter.testelement.ThreadListener;
import org.apache.jmeter.testelement.property.BooleanProperty;
import org.apache.jmeter.threads.JMeterContext;
import org.apache.jmeter.threads.JMeterVariables;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.Serializable;

/**
 * Aes加解密处理器
 * @author guobingbing
 * @version 1.0
 * @wechat t_gbinb
 * @create 2021/6/29 10:23
 */
public class AesSecretProcessor extends AbstractTestElement implements Serializable, PostProcessor, ThreadListener {

	private static final Logger logger = LoggerFactory.getLogger(AesSecretProcessor.class);
	private static final String ALGORITHM = "AES";
	private static final String CIPHER_MODE = "Processor.cipherMode";
	private static final String FILL_MODE = "Processor.fillMode";
	private static final String IV_OFFSET = "Processor.ivOffset";
	private static final String HEX_KEY = "Processor.hexKey";
	private static final String ENCODING = "Processor.encoding";
	private static final String OUT_MODE = "Processor.outMode";
	private static final String INPUT_PARAMS = "Processor.inputParams";
	private static final String OUTPUT_PARAMS = "Processor.outputParams";
	private static final String ENC_OR_DEC = "Processor.encOrDec";
	private static final String USED_PREV_RESULT = "Processor.usedPrevResult";

	/**
	 * ECB/CBC/CTR/OFB/CFB
	 * @param cipherMode
	 */
	public void setCipherMode(String cipherMode) {
		setProperty(CIPHER_MODE, cipherMode.toUpperCase());
	}

	public String getCipherMode(){
		return getPropertyAsString(CIPHER_MODE, "ECB");
	}

	public String getFillMode() {
		return getPropertyAsString(FILL_MODE, "PKCS5Padding");
	}

	/**
	 * PKCS5Padding/PKCS7Padding
	 * @param fillMode
	 */
	public void setFillMode(String fillMode) {
		setProperty(FILL_MODE, fillMode);
	}

	public String getIvOffset() {
		return getPropertyAsString(IV_OFFSET);
	}

	public void setIvOffset(String ivOffset) {
		setProperty(IV_OFFSET, ivOffset);
	}

	public String getHexKey() {
		return getPropertyAsString(HEX_KEY);
	}

	public void setHexKey(String hexKey) {
		setProperty(HEX_KEY, hexKey);
	}

	public String getEncoding() {
		return getPropertyAsString(ENCODING, "utf-8");
	}

	public void setEncoding(String encoding) {
		setProperty(ENCODING, encoding);
	}

	public String getOutMode() {
		return getPropertyAsString(OUT_MODE, "base64");
	}

	public void setOutMode(String outMode) {
		setProperty(OUT_MODE, outMode);
	}

	public void setEncOrDec(String value){
		setProperty(ENC_OR_DEC, value);
	}

	public String getEncOrDec(){
		return getPropertyAsString(ENC_OR_DEC);
	}

	public void setInputParams(String inputParams){
		setProperty(INPUT_PARAMS, inputParams);
	}

	public String getInputParams(){
		return getPropertyAsString(INPUT_PARAMS);
	}

	public void setOutputParams(String outputParams){
		setProperty(OUTPUT_PARAMS, outputParams);
	}

	public String getOutputParams(){
		return getPropertyAsString(OUTPUT_PARAMS);
	}

	public void setUsedPrevResult(boolean value){
		setProperty(new BooleanProperty(USED_PREV_RESULT, value));
	}

	public boolean isUsedPrevResult(){
		return getPropertyAsBoolean(USED_PREV_RESULT);
	}

	@Override
	public void process() {
		JMeterContext ctx = getThreadContext();
		JMeterVariables vars = ctx.getVariables();
		String inputData = "";
		if(isUsedPrevResult()){
			SampleResult prev = ctx.getPreviousResult();
			inputData = prev.getResponseDataAsString();
		}else{
			inputData = vars.get(getInputParams());
		}
		if(StringUtils.isNotBlank(inputData)){
			try{
				String outData = "";
				if("ENC".equalsIgnoreCase(getEncOrDec())){
					outData = encrypt(inputData);
				}else if("DEC".equalsIgnoreCase(getEncOrDec())){
					outData = decrypt(inputData);
				}
				SampleResult prev = ctx.getPreviousResult();
				prev.setResponseData(outData, getEncoding());
				if(StringUtils.isNotBlank(getOutputParams())) {
					vars.put(getOutputParams(), outData);
				}
			}catch(Exception e){
				logger.error("process", e);
			}
		}
	}

	@Override
	public void threadStarted() {
	}

	@Override
	public void threadFinished() {
	}

	private Cipher getCipher(int cipherMode, String encoding) throws Exception{
		SecretKeySpec key = new SecretKeySpec(getHexKey().getBytes(encoding), ALGORITHM);
		String mode = ALGORITHM + "/" + getCipherMode() + "/" + getFillMode();
		Cipher cipher = Cipher.getInstance(mode);
		if(StringUtils.isNotBlank(getIvOffset())){
			IvParameterSpec ivParameterSpec = new IvParameterSpec(getIvOffset().getBytes(encoding));
			cipher.init(cipherMode, key, ivParameterSpec);
		}else{
			cipher.init(cipherMode, key);
		}
		return cipher;
	}

	private String encrypt(String data) throws Exception{
		String encoding = getEncoding();
		Cipher cipher = this.getCipher(Cipher.ENCRYPT_MODE, encoding);
		return Base64.encodeBase64String(cipher.doFinal(data.getBytes(encoding)));
	}

	private String decrypt(String data) throws Exception{
		String encoding = getEncoding();
		Cipher cipher = this.getCipher(Cipher.DECRYPT_MODE, encoding);
		return new String(cipher.doFinal(Base64.decodeBase64(data)));
	}

	public static void main(String[] args) throws Exception {
		AesSecretProcessor processor = new AesSecretProcessor();
		processor.setInputParams("abc");
		processor.setCipherMode("CBC");
		processor.setIvOffset("A-16-Byte-String");
		processor.setHexKey("op09iuj79FD823et");
		processor.setInputParams("wcYKYZsaQysvdPNaxz81oA==");
		System.out.println(processor.getOutputParams());
		System.out.println(processor.encrypt("abc"));
		System.out.println(processor.decrypt("wcYKYZsaQysvdPNaxz81oA=="));
	}
}
